package com.example.bbsbe.controller;

import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

import java.util.Map;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.volcengine.ark.runtime.model.completion.chat.*;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionContentPart.*;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest.*;
import com.volcengine.ark.runtime.service.ArkService;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
/*
# pom.xml
<dependency>
  <groupId>com.volcengine</groupId>
  <artifactId>volcengine-java-sdk-ark-runtime</artifactId>
  // 替换正式版本号
  <version>LATEST</version>
</dependency>
*/


@RestController
@CrossOrigin(origins = "https://yushi.chat", allowCredentials = "true")
public class AiController {

    @GetMapping(value = "/ai_talk/", produces = "text/event-stream")
    public SseEmitter aiTalk(
            @RequestParam String question,
            @RequestParam String maintext,
            @RequestParam String icon1,
            @RequestParam String icon2,
            @RequestParam String icon3) throws JsonProcessingException {

        String apiKey = "b7b04cdd-9df0-41a9-9c82-308c2e81ca81";

        ConnectionPool connectionPool = new ConnectionPool(10, 20, TimeUnit.SECONDS);
        Dispatcher dispatcher = new Dispatcher();
        ArkService service = ArkService.builder()
                .dispatcher(dispatcher)
                .connectionPool(connectionPool)
                .apiKey(apiKey)
                .build();

        SseEmitter emitter = new SseEmitter(60 * 1000L); // 60秒超时
        List<ChatMessage> messagesForReqList = new ArrayList<>();

            List<ChatCompletionContentPart> contentForElementForMessagesForReqList0 = new ArrayList<>();

            ChatCompletionContentPartImageURL
                    imageUrlForElementForContentForElementForMessagesForReqList00 =
                    new ChatCompletionContentPartImageURL();
            imageUrlForElementForContentForElementForMessagesForReqList00.setUrl(
                    icon1);
            ChatCompletionContentPart elementForContentForElementForMessagesForReqList00 =
                    new ChatCompletionContentPart();
            elementForContentForElementForMessagesForReqList00.setType("image_url");
            elementForContentForElementForMessagesForReqList00.setImageUrl(
                    imageUrlForElementForContentForElementForMessagesForReqList00);

        ChatCompletionContentPartImageURL
                imageUrlForElementForContentForElementForMessagesForReqList01 =
                new ChatCompletionContentPartImageURL();
        imageUrlForElementForContentForElementForMessagesForReqList01.setUrl(
                icon2);
        ChatCompletionContentPart elementForContentForElementForMessagesForReqList01 =
                new ChatCompletionContentPart();
        elementForContentForElementForMessagesForReqList01.setType("image_url");
        elementForContentForElementForMessagesForReqList01.setImageUrl(
                imageUrlForElementForContentForElementForMessagesForReqList01);

        ChatCompletionContentPartImageURL
                imageUrlForElementForContentForElementForMessagesForReqList02 =
                new ChatCompletionContentPartImageURL();
        imageUrlForElementForContentForElementForMessagesForReqList02.setUrl(
                icon3);
        ChatCompletionContentPart elementForContentForElementForMessagesForReqList02 =
                new ChatCompletionContentPart();
        elementForContentForElementForMessagesForReqList02.setType("image_url");
        elementForContentForElementForMessagesForReqList02.setImageUrl(
                imageUrlForElementForContentForElementForMessagesForReqList02);

            ChatCompletionContentPart elementForContentForElementForMessagesForReqList03=
                    new ChatCompletionContentPart();
            elementForContentForElementForMessagesForReqList03.setType("text");
            elementForContentForElementForMessagesForReqList03.setText(maintext);
            ChatCompletionContentPart elementForContentForElementForMessagesForReqList04 =
                    new ChatCompletionContentPart();
            elementForContentForElementForMessagesForReqList04.setType("text");
            elementForContentForElementForMessagesForReqList04.setText(question);


            if (!icon1.equals("https://www.yushi.chat/media/0")){
                contentForElementForMessagesForReqList0.add(
                        elementForContentForElementForMessagesForReqList00);
            }
            if (!icon2.equals("https://www.yushi.chat/media/0")){
            contentForElementForMessagesForReqList0.add(
                    elementForContentForElementForMessagesForReqList01);
            }
            if (!icon3.equals("https://www.yushi.chat/media/0")){
            contentForElementForMessagesForReqList0.add(
                    elementForContentForElementForMessagesForReqList02);
            }
            contentForElementForMessagesForReqList0.add(
                    elementForContentForElementForMessagesForReqList03);
            contentForElementForMessagesForReqList0.add(
                    elementForContentForElementForMessagesForReqList04);

            ChatMessage elementForMessagesForReqList0 =
                    ChatMessage.builder()
                            .role(ChatMessageRole.USER)
                            .multiContent(contentForElementForMessagesForReqList0)
                            .build();

            ChatMessage elementForMessagesForReqList1 =
                    ChatMessage.builder()
                            .role(ChatMessageRole.SYSTEM)
                            .content("这是一个动漫论坛，你是一个活泼且了解动漫方面内容的专家小爱，我需要你结合用户所询问的指定帖子内容以及帖子图片回答用户的问题")
                            .build();
            messagesForReqList.add(elementForMessagesForReqList0);
            messagesForReqList.add(elementForMessagesForReqList1);

        ChatCompletionRequest req = ChatCompletionRequest.builder()
                .model("doubao-1-5-thinking-vision-pro-250428")
                .messages(messagesForReqList)
                .stream(true)
                .build();

        // 使用线程池处理流式响应
        service.streamChatCompletion(req)
                .doOnError(error -> {
                    try {
                        emitter.send(SseEmitter.event()
                                .name("error")
                                .data(error.getMessage()));
                        emitter.completeWithError(error);
                    } catch (IOException e) {
                        // 忽略关闭异常
                    }
                })
                .subscribe(
                        choice -> {
                            if (choice.getChoices() != null && !choice.getChoices().isEmpty()) {
                                String content = (String) choice.getChoices().get(0).getMessage().getContent();
                                if (content != null) {
                                    try {
                                        // 发送SSE格式数据
                                        emitter.send(SseEmitter.event()
                                                .name("message")
                                                .data(content));
                                    } catch (IOException e) {
                                        // 处理发送失败
                                    }
                                }
                            }
                        },
                        error -> {
                            // 错误处理
                        },
                        () -> {
                            try {
                                // 发送完成事件
                                emitter.send(SseEmitter.event().name("complete").data("[DONE]"));
                                emitter.complete();
                            } catch (IOException e) {
                                // 忽略关闭异常
                            } finally {
                                service.shutdownExecutor();
                            }
                        }
                );

        // 处理客户端断开连接
        emitter.onCompletion(() -> {
            System.out.println("SSE连接完成");
            service.shutdownExecutor();
        });

        emitter.onTimeout(() -> {
            System.out.println("SSE连接超时");
            emitter.complete();
            service.shutdownExecutor();
        });

        emitter.onError(error -> {
            System.out.println("SSE连接错误: " + error.getMessage());
            service.shutdownExecutor();
        });

        return emitter;
    }
}

